import os
import pickle
import logging
import evaluate
import numpy as np
from transformers import (
    Trainer,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    TrainingArguments,
)
from vllm import LLM, SamplingParams
from modelscope import snapshot_download
from datasets import load_dataset
from dataclasses import dataclass

# from utils import deprecated

import pandas as pd
from datasets import concatenate_datasets, load_from_disk


def save_obj(obj, name):
    """
    将对象保存到文件
    :param obj: 要保存的对象
    :param name: 文件的名称（包括路径）
    """
    with open(name, "wb") as f:
        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)


def load_obj(name):
    """
    从文件加载对象
    :param name: 文件的名称（包括路径）
    :return: 反序列化后的对象
    """
    with open(name, "rb") as f:
        return pickle.load(f)


def glm4_vllm(prompts, output_dir, temperature=0, max_tokens=1024):
    # GLM-4-9B-Chat-1M
    max_model_len, tp_size = 131072, 1
    model_dir = snapshot_download("ZhipuAI/glm-4-9b-chat")

    tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
    llm = LLM(
        model=model_dir,
        tensor_parallel_size=tp_size,
        max_model_len=max_model_len,
        trust_remote_code=True,
        enforce_eager=True,
    )
    stop_token_ids = [151329, 151336, 151338]
    sampling_params = SamplingParams(
        temperature=temperature, max_tokens=max_tokens, stop_token_ids=stop_token_ids
    )

    inputs = tokenizer.apply_chat_template(
        prompts, tokenize=False, add_generation_prompt=True
    )
    outputs = llm.generate(prompts=inputs, sampling_params=sampling_params)

    save_obj(outputs, output_dir)


def get_logger(logger_name, logger_file):
    # 创建Logger
    logger = logging.getLogger(logger_name)
    # logger = logging.getLogger()
    logger.setLevel(logging.DEBUG)
    # 定义handler的输出格式
    formatter = logging.Formatter(
        "%(asctime)s - %(name)s - %(levelname)s: \n%(message)s\n",
        datefmt="%Y/%m/%d %H:%M:%S",
    )
    # 避免重复添加处理器
    if not any(isinstance(h, logging.FileHandler) for h in logger.handlers):
        handler = logging.FileHandler(logger_file)
        handler.setLevel(logging.INFO)

        formatter = logging.Formatter(
            "%(asctime)s - %(name)s - %(levelname)s: \n%(message)s\n",
            datefmt="%Y/%m/%d %H:%M:%S",
        )
        handler.setFormatter(formatter)
        logger.addHandler(handler)

        ch = logging.StreamHandler()
        ch.setLevel(logging.DEBUG)

        ch.setFormatter(formatter)
        logger.addHandler(ch)

    # 写入日志
    # logger.debug("debug message")
    # logger.info("info message")
    # logger.warning("warning message")
    # logger.error("error message")
    # logger.critical("critical message")
    return logger


# export


def get_last_folder_name(path):
    """
    获取文件路径最后一个文件夹的名称
    # 示例用法
    # print(get_last_folder_name('/path/to/your/folder'))
    """
    return os.path.basename(os.path.normpath(path))
